package mongo

import (
	"sync"
	"time"

	"github.com/globalsign/mgo"
	"github.com/globalsign/mgo/bson"
)

var (
	// GDB 定义全局的DB
	GDB  *DB
	once sync.Once
)

// Handler 集合处理函数
type Handler func(c *mgo.Collection) error

// NewDB 实例化mongo数据库
func NewDB(url, dbName string, timeout time.Duration) (*DB, error) {
	if timeout == 0 {
		timeout = 10 * time.Second
	}

	session, err := mgo.DialWithTimeout(url, timeout)
	if err != nil {
		return nil, err
	}

	db := &DB{
		session: session,
		dbName:  dbName,
	}

	once.Do(func() {
		GDB = db
	})
	return db, nil
}

// DB mongo数据库
type DB struct {
	session *mgo.Session
	dbName  string
}

// DBName 数据库名称
func (d *DB) DBName() string {
	return d.dbName
}

// Session 当前会话
func (d *DB) Session() *mgo.Session {
	return d.session
}

// DB 获取当前数据库
func (d *DB) DB(dbName ...string) *mgo.Database {
	name := d.dbName
	if len(dbName) > 0 {
		name = dbName[0]
	}
	return d.Session().DB(name)
}

// CWithDB 获取指定数据库的集合
func (d *DB) CWithDB(dbName, cName string, fn Handler) error {
	session := d.session.Clone()
	defer session.Close()

	return fn(session.DB(dbName).C(cName))
}

// C 集合
func (d *DB) C(cName string, fn Handler) error {
	return d.CWithDB(d.dbName, cName, fn)
}

// IncrID 获取自增ID
func (d *DB) IncrID(cName string) (int64, error) {
	var id int64

	err := d.C("counter", func(c *mgo.Collection) error {
		var result struct {
			Seq int64 `bson:"seq"`
		}
		_, err := c.Find(bson.M{"_id": cName}).Apply(mgo.Change{
			Update:    bson.M{"$inc": bson.M{"seq": 1}},
			ReturnNew: true,
			Upsert:    true,
		}, &result)
		if err != nil {
			return err
		}
		id = result.Seq
		return nil
	})

	if err != nil {
		return id, err
	}
	return id, nil
}

// Close 关闭当前会话
func (d *DB) Close() error {
	d.session.Close()
	return nil
}
